# 导入操作系统相关功能模块，用于文件和目录操作
import os
# 导入中文分词库，用于中文文本分词
import jieba
# 导入高级文件操作模块，用于文件和目录的复制、移动和删除
import shutil
# 导入科学计算库，提供多维数组和矩阵运算功能
import numpy as np
# 从 sklearn 库中导入 TF-IDF 向量化工具，用于文本特征提取
from sklearn.feature_extraction.text import TfidfVectorizer
# 从 sklearn 库中导入 KMeans 聚类算法，用于文本主题聚类
from sklearn.cluster import KMeans
# 导入进度条库，用于显示任务进度
from tqdm import tqdm
# 从 collections 模块中导入计数器类，用于统计元素出现次数
from collections import Counter

# 配置参数
# 输入文件夹路径，指定待处理的文本文件所在目录
input_dir = "E:/jjsw/jjsw-new"  # 输入文件夹路径
# 输出文件夹根目录，指定聚类结果保存的目录
output_dir = "./clusters"  # 输出文件夹根目录
# 聚类数量（主题数量），指定要将文本划分为的主题数量
num_clusters = 6  # 聚类数量（主题数量）
# 处理txt文件的个数
num_files = 1000

# 1. 读取所有txt文件内容
# 该函数用于从指定文件夹中加载所有txt文件内容
# :param folder_path: 包含待处理txt文件的文件夹路径
# :return: 文本内容列表和文件名列表
def load_texts(folder_path):
    # 存储所有文本内容的列表
    texts = []
    # 存储所有文件名的列表
    filenames = []
    # 获取指定文件夹中所有txt文件，并限制只处理前num_files个文件
    txt_files = [filename for filename in os.listdir(folder_path) if filename.endswith(".txt")][:num_files]
    # 遍历所有txt文件，显示加载文件的进度条
    for filename in tqdm(txt_files, desc="加载文件"):
        # 拼接文件的完整路径
        filepath = os.path.join(folder_path, filename)
        # 以只读模式打开文件，并指定编码为utf-8
        with open(filepath, 'r', encoding='utf-8') as f:
            # 读取文件内容
            content = f.read()
            # 将文件内容添加到文本列表中
            texts.append(content)
            # 将文件名添加到文件名列表中
            filenames.append(filename)
    return texts, filenames

# 2. 文本预处理（分词+去停用词）
# 该函数用于对文本进行预处理，包括分词和去停用词操作
# :param texts: 待处理的文本内容列表
# :return: 预处理后的文本内容列表
def preprocess(texts):
    # 加载停用词表
    # 存储停用词的集合
    stopwords = set()
    # 以只读模式打开停用词文件，并指定编码为utf-8
    with open("stopwords.txt", "r", encoding="utf-8") as f:
        # 逐行读取停用词文件
        for line in f:
            # 将每行停用词去除首尾空白后添加到停用词集合中
            stopwords.add(line.strip())
    
    # 存储预处理后文本的列表
    processed_texts = []
    # 遍历所有文本，显示文本预处理的进度条
    for text in tqdm(texts, desc="文本预处理"):
        # 分词
        # 使用 jieba 对文本进行分词
        words = jieba.cut(text)
        # 去停用词和标点
        # 过滤掉停用词和空白词，将剩余词添加到过滤后的词列表中
        filtered_words = [word for word in words if word not in stopwords and word.strip() != '']
        # 将过滤后的词用空格连接成字符串，并添加到预处理后文本列表中
        processed_texts.append(" ".join(filtered_words))
    return processed_texts

# 3. TF-IDF特征提取
# 该函数用于对预处理后的文本进行 TF-IDF 特征提取
# :param texts: 预处理后的文本内容列表
# :return: TF-IDF 矩阵和向量化工具
def vectorize_texts(texts):
    # 初始化 TF-IDF 向量化工具，限制最大特征数量为5000
    vectorizer = TfidfVectorizer(max_features=5000)  # 限制特征数量
    # 对文本进行 TF-IDF 向量化处理
    tfidf_matrix = vectorizer.fit_transform(texts)
    return tfidf_matrix, vectorizer

# 4. 主题聚类
# 该函数用于对 TF-IDF 矩阵进行主题聚类
# :param tfidf_matrix: TF-IDF 特征矩阵
# :param num_clusters: 聚类数量（主题数量）
# :return: 每个文本所属的聚类标签
def cluster_texts(tfidf_matrix, num_clusters):
    # 初始化 KMeans 聚类算法，指定聚类数量和随机种子
    kmeans = KMeans(n_clusters=num_clusters, random_state=42)
    # 使用 TF-IDF 矩阵进行聚类训练
    kmeans.fit(tfidf_matrix)
    return kmeans.labels_

# 5. 创建主题文件夹并移动文件
# 该函数用于提取指定聚类的关键词
# :param texts: 所有文本内容列表
# :param labels: 每个文本所属的聚类标签列表
# :param cluster_id: 要提取关键词的聚类 ID
# :param vectorizer: TF-IDF 向量化工具
# :return: 该聚类的前3个关键词，用下划线连接
def get_cluster_keywords(texts, labels, cluster_id, vectorizer):
    # 筛选出属于指定聚类的文本
    cluster_texts = [texts[i] for i in range(len(texts)) if labels[i] == cluster_id]
    # 对指定聚类的文本进行 TF-IDF 向量化处理
    tfidf_matrix = vectorizer.transform(cluster_texts)
    # 获取特征名称列表
    feature_names = vectorizer.get_feature_names_out()
    # 计算每个特征的 TF-IDF 总和
    tfidf_sum = np.sum(tfidf_matrix, axis=0)
    # 对 TF-IDF 总和进行降序排序，获取排序后的索引
    sorted_indices = np.argsort(tfidf_sum).flatten()[::-1]
    # 获取前3个关键词，并将特征名称转换为字符串
    top_words = [str(feature_names[i]) for i in sorted_indices[:3]]
    return "_".join(top_words)

# 该函数用于根据聚类结果创建主题文件夹，并将文件移动到相应的文件夹中
# :param texts: 所有文本内容列表
# :param filenames: 所有文件名列表
# :param labels: 每个文本所属的聚类标签列表
# :param vectorizer: TF-IDF 向量化工具
def organize_files(texts, filenames, labels, vectorizer):
    # 存储每个聚类的关键词
    cluster_keywords = {}
    # 遍历所有聚类 ID
    for cluster_id in set(labels):
        # 提取指定聚类的关键词
        keywords = get_cluster_keywords(texts, labels, cluster_id, vectorizer)
        # 将聚类 ID 和对应的关键词添加到字典中
        cluster_keywords[cluster_id] = keywords
    
    # 遍历所有文本的聚类标签，显示文件归类的进度条
    for idx, label in tqdm(enumerate(labels), desc="文件归类", total=len(labels)):
        # 拼接聚类文件夹的完整路径
        cluster_folder = os.path.join(output_dir, f"cluster_{cluster_keywords[label]}")
        # 创建聚类文件夹，如果文件夹已存在则不报错
        os.makedirs(cluster_folder, exist_ok=True)
        # 拼接源文件的完整路径
        src_file = os.path.join(input_dir, filenames[idx])
        # 拼接目标文件的完整路径
        dst_file = os.path.join(cluster_folder, filenames[idx])
        # 将源文件移动到目标文件夹中
        shutil.move(src_file, dst_file)

if __name__ == "__main__":
    # 清空 clusters 文件夹
    # 检查输出文件夹是否存在
    if os.path.exists(output_dir):
        # 如果存在，则删除该文件夹及其所有内容
        shutil.rmtree(output_dir)
    # 创建输出文件夹，如果文件夹已存在则不报错
    os.makedirs(output_dir, exist_ok=True)

    # 执行流程
    # 加载文本文件内容和文件名
    texts, filenames = load_texts(input_dir)
    # 打印加载的文本文件数量
    print(f"共加载 {len(texts)} 个文本文件...")
    
    # 对文本进行预处理
    processed_texts = preprocess(texts)
    # 打印文本预处理完成信息
    print("完成文本预处理...")
    
    # 对预处理后的文本进行 TF-IDF 特征提取
    tfidf_matrix, vectorizer = vectorize_texts(processed_texts)
    # 打印 TF-IDF 特征提取完成信息
    print("完成TF-IDF特征提取...")
    
    # 对 TF-IDF 矩阵进行主题聚类
    labels = cluster_texts(tfidf_matrix, num_clusters)
    # 打印聚类完成信息和发现的主题数量
    print(f"完成聚类，共发现 {num_clusters} 个主题...")
    
    # 根据聚类结果对文件进行归类
    organize_files(texts, filenames, labels, vectorizer)
    # 打印文件归类完成信息
    print("文件归类完成！")